{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "f1f73c8e-62a1-4431-b69a-3e0217e4256f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1cd0371e0eec43a497c455b1d81b9ee7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Generating train split: 0 examples [00:00, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "ds = load_dataset('json', data_files='/root/autodl-tmp/dataset/self_recognition/self_cognition.jsonl')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "e92e5cfa-4edd-4929-b1db-f76d447f809e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'query': '你是？',\n",
       " 'response': '我是{{NAME}}，由{{AUTHOR}}训练的人工智能助手。我的目标是为用户提供有用、准确和及时的信息，并通过各种方式帮助用户进行有效的沟通。请告诉我有什么可以帮助您的呢？',\n",
       " 'tag': 'zh'}"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ds['train'][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "96f7009f-63c1-45df-bf7a-2a5baf9249df",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "成功转换 108 条数据\n",
      "{\n",
      "  \"system\": \"你是一个有帮助的AI助手。\",\n",
      "  \"messages\": [\n",
      "    {\n",
      "      \"from\": \"human\",\n",
      "      \"value\": \"你是？\"\n",
      "    },\n",
      "    {\n",
      "      \"from\": \"gpt\",\n",
      "      \"value\": \"我是天天，由中国移动训练的人工智能助手。我的目标是为用户提供有用、准确和及时的信息，并通过各种方式帮助用户进行有效的沟通。请告诉我有什么可以帮助您的呢？\"\n",
      "    }\n",
      "  ]\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "import json\n",
    "\n",
    "def convert_to_sharegpt_format(data, name='天天', author='中国移动',\n",
    "                               system_prompt=\"你是一个有帮助的AI助手。\"):\n",
    "    \"\"\"\n",
    "    将包含多轮对话的数据转换为ShareGPT格式\n",
    "    \n",
    "    参数:\n",
    "        data: 原始数据列表，每个元素包含query和response字段\n",
    "        name: 替换{{NAME}}占位符的值\n",
    "        author: 替换{{AUTHOR}}占位符的值\n",
    "        system_prompt: 系统提示词，默认为通用助手描述\n",
    "    \"\"\"\n",
    "    sharegpt_data = []\n",
    "    \n",
    "    for item in data:\n",
    "        # 构建ShareGPT格式\n",
    "        sharegpt_item = {\n",
    "            \"system\": system_prompt,\n",
    "            \"messages\": []\n",
    "        }\n",
    "        \n",
    "        # 获取查询和回复\n",
    "        query = item.get(\"query\", \"\")\n",
    "        response = item.get(\"response\", \"\").strip()\n",
    "        \n",
    "        # 替换回复中的占位符\n",
    "        response = response.replace(\"{{NAME}}\", name).replace(\"{{AUTHOR}}\", author)\n",
    "        \n",
    "        # 添加用户消息\n",
    "        sharegpt_item[\"messages\"].append({\n",
    "            \"from\": \"human\",\n",
    "            \"value\": query\n",
    "        })\n",
    "        \n",
    "        # 添加AI回复\n",
    "        sharegpt_item[\"messages\"].append({\n",
    "            \"from\": \"gpt\",\n",
    "            \"value\": response\n",
    "        })\n",
    "        \n",
    "        sharegpt_data.append(sharegpt_item)\n",
    "    \n",
    "    return sharegpt_data\n",
    "\n",
    "\n",
    "# 转换数据\n",
    "sharegpt_data = convert_to_sharegpt_format(ds['train'])\n",
    "\n",
    "# 输出结果\n",
    "if sharegpt_data:\n",
    "    print(f\"成功转换 {len(sharegpt_data)} 条数据\")\n",
    "    # 打印第一条示例\n",
    "    print(json.dumps(sharegpt_data[0], ensure_ascii=False, indent=2))\n",
    "else:\n",
    "    print(\"转换失败，未识别到有效对话\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "075d2867-3d26-41f0-a3ce-6815c4fdf67a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "使用特殊标识符的结构化文本：\n",
      "\n",
      "\n",
      "<|im_start|>system\n",
      "你是一个有帮助的AI助手。\n",
      "<|im_start|>human\n",
      "你是谁!\n",
      "<|im_start|>assistant\n",
      "您好！我是中国移动开发的人工智能语言模型，名为天天。我可以回答您的问题、提供信息、进行对话并帮助解决问题。如果您有任何疑问或需要帮助，请随时告诉我！<|im_end|>\n"
     ]
    }
   ],
   "source": [
    "from typing import List, Dict\n",
    "\n",
    "def convert_sharegpt_to_formatted_text(sharegpt_data: List[Dict]):\n",
    "    \"\"\"\n",
    "    将ShareGPT格式数据拼接为结构化文本序列\n",
    "    \n",
    "    参数:\n",
    "        sharegpt_data: ShareGPT格式数据列表\n",
    "    \"\"\"\n",
    "    formatted_texts = []\n",
    "    \n",
    "    # 定义特殊标识符\n",
    "    im_start = \"<|im_start|>\"\n",
    "    im_end = \"<|im_end|>\"\n",
    "    nl_tokens = \"\\n\"\n",
    "    _system = f\"system{nl_tokens}\"\n",
    "    _user = f\"human{nl_tokens}\"\n",
    "    _assistant = f\"assistant{nl_tokens}\"\n",
    "    \n",
    "    for item in sharegpt_data:\n",
    "        system_prompt = item.get(\"system\", \"\")\n",
    "        messages = item.get(\"messages\", [])\n",
    "        \n",
    "        if not messages:\n",
    "            continue\n",
    "            \n",
    "        # 拼接结构化文本\n",
    "        formatted_text = \"\"\n",
    "        \n",
    "        # 添加系统提示词\n",
    "        if system_prompt:\n",
    "            formatted_text += f\"{im_start}{_system}{system_prompt}{nl_tokens}\"\n",
    "        \n",
    "        # 拼接多轮对话\n",
    "        for msg in messages:\n",
    "            role = msg.get(\"from\")\n",
    "            content = msg.get(\"value\", \"\").strip()\n",
    "            \n",
    "            if not role or not content:\n",
    "                continue\n",
    "                \n",
    "            # 根据角色添加标识符\n",
    "            if role == \"human\":\n",
    "                role_identifier = f\"{im_start}{_user}\"\n",
    "            elif role == \"gpt\":\n",
    "                role_identifier = f\"{im_start}{_assistant}\"\n",
    "            else:\n",
    "                role_identifier = f\"{im_start}{role}{nl_tokens}\"\n",
    "            \n",
    "            # 拼接角色和内容\n",
    "            formatted_text += f\"{role_identifier}{content}\"\n",
    "            \n",
    "            # 助手回答后添加结束标记\n",
    "            if role == \"gpt\":\n",
    "                formatted_text += f\"{im_end}{nl_tokens}\"\n",
    "            else:\n",
    "                formatted_text += f\"{nl_tokens}\"\n",
    "        \n",
    "        formatted_texts.append(formatted_text.strip())\n",
    "    \n",
    "    return formatted_texts\n",
    "\n",
    "\n",
    "# 拼接为结构化文本\n",
    "formatted_text = convert_sharegpt_to_formatted_text(sharegpt_data)\n",
    "print(\"使用特殊标识符的结构化文本：\\n\\n\")\n",
    "print(formatted_text[1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "16ec8188-8ead-49e4-bfd8-66e9524efd27",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from typing import List, Dict, Any\n",
    "\n",
    "\n",
    "def process_texts_to_dataset(\n",
    "    text_list: List[str],\n",
    "    tokenizer,\n",
    "    max_length: int\n",
    ") -> Dict[str, torch.Tensor]:\n",
    "    \"\"\"\n",
    "    将结构化文本列表转换为模型输入数据（input_ids/labels/attention_mask）\n",
    "    \n",
    "    参数:\n",
    "        text_list: 结构化文本列表（每条文本已用<|im_start|>/<|im_end|>标记）\n",
    "        tokenizer: 预训练Tokenizer实例\n",
    "        max_length: 最大序列长度（截断&填充目标）\n",
    "    \"\"\"\n",
    "    # 处理pad_token（若未定义则使用eos_token）\n",
    "    if tokenizer.pad_token is None:\n",
    "        tokenizer.pad_token = tokenizer.eos_token\n",
    "    \n",
    "    all_input_ids = []\n",
    "    all_labels = []\n",
    "    all_attention_mask = []\n",
    "    \n",
    "    # 标记assistant内容的开始/结束标识符\n",
    "    ASSISTANT_START = \"<|im_start|>assistant\\n\"\n",
    "    ASSISTANT_END = \"<|im_end|>\"\n",
    "    \n",
    "    # 解析标识符的Token ID\n",
    "    start_tokens = tokenizer.encode(ASSISTANT_START, add_special_tokens=False)\n",
    "    end_token = tokenizer.encode(ASSISTANT_END, add_special_tokens=False)[0]\n",
    "    \n",
    "    for text in text_list:\n",
    "        # 1. Tokenize文本并截断\n",
    "        tokens = tokenizer.encode(text, max_length=max_length, truncation=True)\n",
    "        input_ids = tokens.copy()\n",
    "        \n",
    "        # 2. 生成attention_mask\n",
    "        attention_mask = [1] * len(input_ids)\n",
    "        \n",
    "        # 3. 初始化labels（全部设为-100，后续只保留assistant内容）\n",
    "        labels = [-100] * len(input_ids)\n",
    "        \n",
    "        # 4. 定位assistant内容并标记labels\n",
    "        i = 0\n",
    "        while i < len(input_ids) - len(start_tokens):\n",
    "            # 检测assistant块开始\n",
    "            if input_ids[i:i+len(start_tokens)] == start_tokens:\n",
    "                start_pos = i + len(start_tokens)\n",
    "                i = start_pos  # 跳过标识符\n",
    "                \n",
    "                # 寻找assistant块结束\n",
    "                end_pos = len(input_ids)\n",
    "                for j in range(start_pos, len(input_ids)):\n",
    "                    if input_ids[j] == end_token:\n",
    "                        end_pos = j\n",
    "                        break\n",
    "                \n",
    "                # 将assistant内容设为有效labels\n",
    "                if start_pos < end_pos:\n",
    "                    labels[start_pos:end_pos] = input_ids[start_pos:end_pos]\n",
    "                i = end_pos + 1  # 跳过结束符\n",
    "            else:\n",
    "                i += 1\n",
    "        \n",
    "        # 5. 过滤无效数据（无assistant内容）\n",
    "        if all(l == -100 for l in labels):\n",
    "            continue\n",
    "        \n",
    "        # 6. 保存处理结果\n",
    "        all_input_ids.append(input_ids)\n",
    "        all_labels.append(labels)\n",
    "        all_attention_mask.append(attention_mask)\n",
    "    \n",
    "    # 7. 批量填充到统一长度\n",
    "    if not all_input_ids:\n",
    "        raise ValueError(\"输入数据中无有效assistant内容，无法创建数据集\")\n",
    "    \n",
    "    max_len = max(len(ids) for ids in all_input_ids)\n",
    "    pad_id = tokenizer.pad_token_id\n",
    "    \n",
    "    # 初始化张量\n",
    "    input_ids_tensor = torch.full((len(all_input_ids), max_len), pad_id, dtype=torch.long)\n",
    "    labels_tensor = torch.full((len(all_labels), max_len), -100, dtype=torch.long)\n",
    "    attention_mask_tensor = torch.zeros((len(all_attention_mask), max_len), dtype=torch.long)\n",
    "    \n",
    "    # 填充数据\n",
    "    for i, (ids, labs, mask) in enumerate(zip(all_input_ids, all_labels, all_attention_mask)):\n",
    "        input_ids_tensor[i, :len(ids)] = torch.tensor(ids)\n",
    "        labels_tensor[i, :len(labs)] = torch.tensor(labs)\n",
    "        attention_mask_tensor[i, :len(mask)] = torch.tensor(mask)\n",
    "    \n",
    "    return {\n",
    "        \"input_ids\": input_ids_tensor,\n",
    "        \"labels\": labels_tensor,\n",
    "        \"attention_mask\": attention_mask_tensor\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "054945b5-ea92-42cf-aa7c-ca9de5df4732",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "78b3fe7a-8b45-4f43-8688-0dc64b27f4fb",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/root/miniconda3/lib/python3.8/site-packages/torchvision/io/image.py:13: UserWarning: Failed to load image Python extension: libtorch_cuda_cu.so: cannot open shared object file: No such file or directory\n",
      "  warn(f\"Failed to load image Python extension: {e}\")\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import json\n",
    "from tqdm import tqdm\n",
    "\n",
    "# 用于将数据处理为torch张量形式\n",
    "import torch\n",
    "from torch.utils.data import Dataset\n",
    "\n",
    "# 用于加载大模型\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
    "\n",
    "# 设置大模型在本地的路径\n",
    "model_path = \"/root/autodl-tmp/qwen2_05b\"\n",
    "\n",
    "# 加载分词器和模型\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_path)\n",
    "\n",
    "model = AutoModelForCausalLM.from_pretrained(model_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "2ece711a-0b67-4f6f-9eb4-2fedad848e0c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def preprocess(dataset, tokenizer, max_length):\n",
    "    sharegpt_data = convert_to_sharegpt_format(dataset)\n",
    "    formatted_texts = convert_sharegpt_to_formatted_text(sharegpt_data)\n",
    "    result = process_texts_to_dataset(formatted_texts, tokenizer, max_length)\n",
    "    return result\n",
    "    \n",
    "\n",
    "ds = load_dataset('json', data_files='/root/autodl-tmp/dataset/self_recognition/self_cognition.jsonl')\n",
    "ds = ds['train']\n",
    "\n",
    "result = preprocess(ds, tokenizer, 1024)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "1c0d36dd-9b12-4c97-828e-3c0fa08a824c",
   "metadata": {},
   "outputs": [],
   "source": [
    "class SupervisedDataset(Dataset):\n",
    "\n",
    "    def __init__(self, dataset, tokenizer, max_len: int):\n",
    "        super(SupervisedDataset, self).__init__()\n",
    "\n",
    "        data_dict = preprocess(dataset, tokenizer, max_len)\n",
    "\n",
    "        self.input_ids = data_dict[\"input_ids\"]\n",
    "        self.labels = data_dict[\"labels\"]\n",
    "        self.attention_mask = data_dict[\"attention_mask\"]\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.input_ids)\n",
    "\n",
    "    def __getitem__(self, i) -> Dict[str, torch.Tensor]:\n",
    "        return dict(\n",
    "            input_ids=self.input_ids[i],\n",
    "            labels=self.labels[i],\n",
    "            attention_mask=self.attention_mask[i],\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "66e6c139-1e02-41d1-8ac8-08229884d78f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "from transformers import (\n",
    "    AutoModelForCausalLM,\n",
    "    AutoTokenizer,\n",
    "    TrainingArguments,\n",
    "    Trainer,\n",
    "    DataCollatorForLanguageModeling\n",
    ")\n",
    "from peft import LoraConfig, get_peft_model\n",
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "4db35352-54c5-42d7-ade4-0f0112bb9aff",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "trainable params: 1,081,344 || all params: 495,114,112 || trainable%: 0.2184\n"
     ]
    }
   ],
   "source": [
    "# 配置LoRA参数\n",
    "lora_config = LoraConfig(\n",
    "    r=16,                  # LoRA注意力维度\n",
    "    lora_alpha=32,         # Alpha参数\n",
    "    target_modules=[\"q_proj\", \"v_proj\"],  # 对哪些模块应用LoRA\n",
    "    lora_dropout=0.2,     # Dropout概率\n",
    "    bias=\"none\",           # 是否训练偏置\n",
    "    task_type=\"CAUSAL_LM\"  # 任务类型\n",
    ")\n",
    "\n",
    "# 将LoRA配置应用到模型\n",
    "model = get_peft_model(model, lora_config)\n",
    "model.print_trainable_parameters()  # 打印可训练参数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "6262e5f4-206a-435e-9d01-a76af57ab7ca",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='5400' max='5400' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [5400/5400 06:24, Epoch 200/200]\n",
       "    </div>\n",
       "    <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       " <tr style=\"text-align: left;\">\n",
       "      <th>Step</th>\n",
       "      <th>Training Loss</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>100</td>\n",
       "      <td>1.894900</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>200</td>\n",
       "      <td>1.218800</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>300</td>\n",
       "      <td>0.974900</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>400</td>\n",
       "      <td>0.783700</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>500</td>\n",
       "      <td>0.654100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>600</td>\n",
       "      <td>0.540500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>700</td>\n",
       "      <td>0.422300</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>800</td>\n",
       "      <td>0.324900</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>900</td>\n",
       "      <td>0.254700</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1000</td>\n",
       "      <td>0.214700</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1100</td>\n",
       "      <td>0.162000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1200</td>\n",
       "      <td>0.126000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1300</td>\n",
       "      <td>0.100200</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1400</td>\n",
       "      <td>0.081800</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1500</td>\n",
       "      <td>0.069400</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1600</td>\n",
       "      <td>0.055900</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1700</td>\n",
       "      <td>0.048300</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1800</td>\n",
       "      <td>0.040200</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1900</td>\n",
       "      <td>0.034700</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2000</td>\n",
       "      <td>0.028400</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2100</td>\n",
       "      <td>0.029600</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2200</td>\n",
       "      <td>0.023100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2300</td>\n",
       "      <td>0.021000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2400</td>\n",
       "      <td>0.019900</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2500</td>\n",
       "      <td>0.017900</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2600</td>\n",
       "      <td>0.014000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2700</td>\n",
       "      <td>0.013000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2800</td>\n",
       "      <td>0.011900</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2900</td>\n",
       "      <td>0.010700</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3000</td>\n",
       "      <td>0.011900</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3100</td>\n",
       "      <td>0.010500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3200</td>\n",
       "      <td>0.010400</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3300</td>\n",
       "      <td>0.009800</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3400</td>\n",
       "      <td>0.010300</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3500</td>\n",
       "      <td>0.008500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3600</td>\n",
       "      <td>0.008200</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3700</td>\n",
       "      <td>0.010200</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3800</td>\n",
       "      <td>0.008500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3900</td>\n",
       "      <td>0.007600</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>4000</td>\n",
       "      <td>0.006700</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>4100</td>\n",
       "      <td>0.007000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>4200</td>\n",
       "      <td>0.006000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>4300</td>\n",
       "      <td>0.006700</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>4400</td>\n",
       "      <td>0.006600</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>4500</td>\n",
       "      <td>0.006300</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>4600</td>\n",
       "      <td>0.006000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>4700</td>\n",
       "      <td>0.005900</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>4800</td>\n",
       "      <td>0.005000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>4900</td>\n",
       "      <td>0.005400</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>5000</td>\n",
       "      <td>0.005400</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>5100</td>\n",
       "      <td>0.005100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>5200</td>\n",
       "      <td>0.005100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>5300</td>\n",
       "      <td>0.005000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>5400</td>\n",
       "      <td>0.005000</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "TrainOutput(global_step=5400, training_loss=0.15508673288204053, metrics={'train_runtime': 384.5131, 'train_samples_per_second': 56.175, 'train_steps_per_second': 14.044, 'total_flos': 5955038399692800.0, 'train_loss': 0.15508673288204053, 'epoch': 200.0})"
      ]
     },
     "execution_count": 41,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_dataset = SupervisedDataset(ds, tokenizer, max_len=512)\n",
    "\n",
    "# 设置训练参数\n",
    "training_args = TrainingArguments(\n",
    "    output_dir=\"./results\",\n",
    "    overwrite_output_dir=True,\n",
    "    num_train_epochs=200,\n",
    "    per_device_train_batch_size=4,\n",
    "    save_steps=10_000,\n",
    "    save_total_limit=2,\n",
    "    prediction_loss_only=True,\n",
    "    logging_steps=100,\n",
    "    save_strategy=\"no\",  # 禁用自动保存检查点\n",
    ")\n",
    "\n",
    "# 创建Trainer\n",
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    args=training_args,\n",
    "    train_dataset=train_dataset,\n",
    ")\n",
    "\n",
    "# 开始训练\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "b7b02720-a825-4dd8-90ef-d42bc1da682a",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:None for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "system\n",
      "你是一个有帮助的AI助手。\n",
      "human\n",
      "介绍一下自己？\n",
      "assistant 天天是中国移动创建的人工智能对话员和对话生成系统。它能够与人类进行对话并回答问题，以便更好地与他们进行交流。天天能够回答各种问题，并且能够提供有用的信息。天天也可以回答问题并提供有用的信息。天天能够回答问题并提供有用的信息。天天能够回答问题并提供有用的信息。天天能够回答问题并提供有用的信息。天天能够回答问题并提供有用的信息。天天能够回答\n"
     ]
    }
   ],
   "source": [
    "# 训练完成后，直接使用model对象进行推理\n",
    "model.eval()  # 设置为评估模式\n",
    "\n",
    "# 准备输入\n",
    "prompt = \"\"\"<|im_start|>system\n",
    "你是一个有帮助的AI助手。\n",
    "<|im_start|>human\n",
    "介绍一下自己？\n",
    "<|im_start|>assistant\"\"\"\n",
    "inputs = tokenizer(prompt, return_tensors=\"pt\").to(model.device)\n",
    "\n",
    "# 生成文本\n",
    "with torch.no_grad():  # 推理时不需要梯度\n",
    "    outputs = model.generate(\n",
    "        **inputs,\n",
    "        max_new_tokens=100,\n",
    "        temperature=0.1,\n",
    "        do_sample=True\n",
    "    )\n",
    "\n",
    "# 解码生成的文本\n",
    "generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)\n",
    "print(generated_text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "fbebd7f2-dc5d-4ca5-99ba-7661c4125d5f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "PeftModelForCausalLM(\n",
       "  (base_model): LoraModel(\n",
       "    (model): PeftModelForCausalLM(\n",
       "      (base_model): LoraModel(\n",
       "        (model): Qwen2ForCausalLM(\n",
       "          (model): Qwen2Model(\n",
       "            (embed_tokens): Embedding(151936, 896)\n",
       "            (layers): ModuleList(\n",
       "              (0-23): 24 x Qwen2DecoderLayer(\n",
       "                (self_attn): Qwen2SdpaAttention(\n",
       "                  (q_proj): lora.Linear(\n",
       "                    (base_layer): Linear(in_features=896, out_features=896, bias=True)\n",
       "                    (lora_dropout): ModuleDict(\n",
       "                      (default): Dropout(p=0.05, inplace=False)\n",
       "                    )\n",
       "                    (lora_A): ModuleDict(\n",
       "                      (default): Linear(in_features=896, out_features=16, bias=False)\n",
       "                    )\n",
       "                    (lora_B): ModuleDict(\n",
       "                      (default): Linear(in_features=16, out_features=896, bias=False)\n",
       "                    )\n",
       "                    (lora_embedding_A): ParameterDict()\n",
       "                    (lora_embedding_B): ParameterDict()\n",
       "                    (lora_magnitude_vector): ModuleDict()\n",
       "                  )\n",
       "                  (k_proj): Linear(in_features=896, out_features=128, bias=True)\n",
       "                  (v_proj): lora.Linear(\n",
       "                    (base_layer): Linear(in_features=896, out_features=128, bias=True)\n",
       "                    (lora_dropout): ModuleDict(\n",
       "                      (default): Dropout(p=0.05, inplace=False)\n",
       "                    )\n",
       "                    (lora_A): ModuleDict(\n",
       "                      (default): Linear(in_features=896, out_features=16, bias=False)\n",
       "                    )\n",
       "                    (lora_B): ModuleDict(\n",
       "                      (default): Linear(in_features=16, out_features=128, bias=False)\n",
       "                    )\n",
       "                    (lora_embedding_A): ParameterDict()\n",
       "                    (lora_embedding_B): ParameterDict()\n",
       "                    (lora_magnitude_vector): ModuleDict()\n",
       "                  )\n",
       "                  (o_proj): Linear(in_features=896, out_features=896, bias=False)\n",
       "                  (rotary_emb): Qwen2RotaryEmbedding()\n",
       "                )\n",
       "                (mlp): Qwen2MLP(\n",
       "                  (gate_proj): Linear(in_features=896, out_features=4864, bias=False)\n",
       "                  (up_proj): Linear(in_features=896, out_features=4864, bias=False)\n",
       "                  (down_proj): Linear(in_features=4864, out_features=896, bias=False)\n",
       "                  (act_fn): SiLU()\n",
       "                )\n",
       "                (input_layernorm): Qwen2RMSNorm((896,), eps=1e-06)\n",
       "                (post_attention_layernorm): Qwen2RMSNorm((896,), eps=1e-06)\n",
       "              )\n",
       "            )\n",
       "            (norm): Qwen2RMSNorm((896,), eps=1e-06)\n",
       "            (rotary_emb): Qwen2RotaryEmbedding()\n",
       "          )\n",
       "          (lm_head): Linear(in_features=896, out_features=151936, bias=False)\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e0b5eeb-8eac-448e-a653-58478c368dcf",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
